package com.log.enhance.thread;

import com.log.enhance.util.ThreadMdcUtil;
import org.slf4j.MDC;
import org.springframework.lang.NonNull;

import java.util.concurrent.Callable;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.ForkJoinTask;

/**
 * 包装ForkJoinPool线程池
 *
 * @author Mr_wenpan@163.com 2021/06/01 15:49
 */
public class MdcWrapperForkJoinPool extends ForkJoinPool {

    public MdcWrapperForkJoinPool() {
        super();
    }

    public MdcWrapperForkJoinPool(int parallelism) {
        super(parallelism);
    }

    public MdcWrapperForkJoinPool(int parallelism, ForkJoinWorkerThreadFactory factory,
                                  Thread.UncaughtExceptionHandler handler, boolean asyncMode) {
        super(parallelism, factory, handler, asyncMode);
    }

    /**
     * 包装父类的执行任务的方法
     *
     * @param task 任务
     * @author Mr_wenpan@163.com 2021/6/1 3:50 下午
     */
    @Override
    public void execute(Runnable task) {
        // MDC.getCopyOfContextMap()：线程提交任务时，获取当前线程中的MDC中的数据
        super.execute(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }

    @Override
    @NonNull
    public <T> ForkJoinTask<T> submit(Runnable task, T result) {
        return super.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()), result);
    }

    @Override
    @NonNull
    public <T> ForkJoinTask<T> submit(Callable<T> task) {
        return super.submit(ThreadMdcUtil.wrap(task, MDC.getCopyOfContextMap()));
    }
}
